from flask import Flask, render_template, request, redirect, url_for,send_from_directory,send_file
from flask import session,jsonify
import pygame._sdl2 as sdl2
from pygame import mixer
import pygame.mixer
import keyboard
import threading
import datetime
import requests
import socket
import ollama
import random
import shutil
import wave
import json
import time
import os 
import re

app = Flask(__name__)

app.secret_key = ''

from volcenginesdkarkruntime import Ark

client = Ark(
    api_key = 'a43dca4d-03e2-40c1-bb77-a149caf5416a', 
    timeout = 1800,
)

# 打印分割线
def print_split_line():
    width = os.get_terminal_size().columns
    text = '=' * width
    print(text)


# 平台列表
platforms = []

for platform in os.listdir(os.path.join('config','话术脚本')):
    if platform.startswith('.'):
        continue

    platforms.append(platform)

# 产品列表
products = []

for platform in os.listdir(os.path.join('config','话术脚本')):
    if platform.startswith('.'):
        continue

    for product in os.listdir(os.path.join('config','话术脚本',platform)):
        if product.startswith('.'):
            continue

        products.append(product)

# 音色列表
streamers = [streamer for streamer in os.listdir(os.path.join('D:\\','音色定制')) if os.path.isdir(os.path.join('D:\\','音色定制',streamer))]

task_config = {}

lock = threading.Lock()

# 历史文案加载
hist_script_content = {}

for platform in os.listdir('D:\\历史文案'):
    if not os.path.isdir(os.path.join('D:\\历史文案',platform)):
        continue

    hist_script_content[platform] = ""

    for file in os.listdir(os.path.join('D:\\历史文案',platform)):
        if not file.endswith('.txt'):
            continue

        with open(os.path.join('D:\\历史文案',platform,file),'r+',encoding = 'utf-8') as f:
            hist_script_content[platform] += re.sub('[，|。|！|!|？|?|:|：|\n|-|——|;|；|、]','',f.read().strip()) 


    print('加载%s历史文案%d份' % (platform,len(os.listdir(os.path.join('D:\\历史文案',platform)))))
    print_split_line()


# 库存文案加载
for platform in os.listdir(os.path.join('config','话术脚本')):
    if not os.path.isdir(os.path.join('config','话术脚本',platform)):
        continue

    for product_folder in os.listdir(os.path.join('config','话术脚本',platform)):
        if not os.path.isdir(os.path.join('config','话术脚本',platform,product_folder)):
            continue

        if os.path.exists(os.path.join('config','话术脚本',platform,product_folder,'文本库')):
            for streamer_folder in os.listdir(os.path.join('config','话术脚本',platform,product_folder,'文本库')):
                if not os.path.isdir(os.path.join('config','话术脚本',platform,product_folder,'文本库',streamer_folder)):
                    continue

                for file in os.listdir(os.path.join('config','话术脚本',platform,product_folder,'文本库',streamer_folder)):
                    if not file.endswith('.txt'):
                        continue

                    with open(os.path.join('config','话术脚本',platform,product_folder,'文本库',streamer_folder,file),'r+',encoding = 'utf-8') as f:
                        hist_script_content[platform] += re.sub('[，|。|！|!|？|?|:|：|\n|-|——|;|；|、]','',f.read().strip()) 

                print(platform,product_folder,streamer_folder,'加载库存文案',len(os.listdir(os.path.join('config','话术脚本',platform,product_folder,'文本库',streamer_folder))),'份')


for scene in ['scene-1','scene-2','scene-3','scene-4','scene-5','scene-6','scene-7','scene-8','scene-9','scene-10']:
    task_config[scene] = {
        'scene_name': scene,
        'current_product': '',
        'current_streamer':'',
        'live_platform':'',
        'startGenScript': False,
        'cache_count':'10',
        'active': True if scene == 'scene-1' else False
    }


def qwen3(prompt,temperature = None):
    if temperature:
        resp = ollama.generate(
            model = 'qwen3:1.7b',
            prompt = prompt,
            options = {
                'temperature':temperature
            }
        )
    else:
        resp = ollama.generate(
            model = 'qwen3:1.7b',
            prompt = prompt,
            options = {

            }
        )
    
    resp['response'] = re.sub('（.+?）|【.+?】|\(.+?\)','',resp['response'],flags = re.DOTALL)
    resp['response'] = re.sub(r'<think>(.*?)</think>','',resp['response'],flags = re.DOTALL)
    resp['response'] = re.sub(r'</think>.*','',resp['response'],flags = re.DOTALL)
    resp['response'] = re.sub('\n\n','\n',resp['response'])
    resp['response'] = re.sub('\*|\-|#','',resp['response'])

    return resp['response']


# 将prompt随机打乱
def random_prompt(product_prompt):
    prompt_lines = [i for i in re.findall('●.+\n',product_prompt)]
    
    random.shuffle(prompt_lines)
    
    # 提取【背景】
    background = [i + '\n' for i in re.findall('【背景】.+。',product_prompt)]

    # 提取【要求】
    requests = [i + '\n' for i in re.findall('【要求】.+。',product_prompt)]
    
    return ''.join(background) + ''.join(prompt_lines) + ''.join(requests)


def doubao(content,temperature = 0.8):
    print_split_line()
    print('提示词：\n',content)
    print('model：doubao-seed-1.6-250615')
    print('temperature：',temperature)
    print('正在请求豆包...')

    response = client.chat.completions.create(
        model = "doubao-seed-1.6-250615",
        messages = [
            {
                "role": "user",
                "content": content
            }
        ],
        thinking = {
             "type": "enabled" 
        },
        temperature = temperature
    )
    
    return response.choices[0].message.content


# 清洗话术
def clear_script(content,blocked_words,live_platform):
    global hist_script_content

    if not live_platform in hist_script_content:
        hist_script_content[live_platform] = ''

    for i in range(10):
        content = re.sub('\*|……|✅| ','',content)
        content = re.sub('【.+?】|（.+?）','，',content)
        content = re.sub('###.+\n','',content)
        content = re.sub('[一二三四五六七八九十]、','',content)
        content = re.sub('~','。',content)
        content = re.sub('\n\n','\n',content)
        content = re.sub(' +\n','\n',content)
        content = re.sub('^[1-9]\.','',content)

    result = []

    for line in content.split():
        if not line.strip():
            continue

        if len(line) < 30:
            continue

        line_ = re.sub('[，|。|！|!|？|?|:|：|\n|-|——|;|；|、]','',line) 
        
        invalide = False

        # 查询句子是否超过30个字在历史上重复过
        for start in range(0,len(line_) - 30):
            if line_[start:start + 30] in hist_script_content[live_platform]:
                print(line_[start:start + 30],'在历史上出现过')
                print_split_line()
                invalide = True
                break

        # 句子内部在重复
        for start in range(0,len(line_) - 10):
            if line_.count(line_[start:start + 10]) > 1:
                print(line_[start:start + 10],'在内部出现过一次以上')
                print_split_line()
                invalide = True
                break

        # 和本次生成的整个文件进行去重
        for start in range(0,len(line_) - 20):
            for li in result:
                if line_[start:start + 20] in re.sub('[，|。|！|!|？|?|:|：|\n|-|——|;|；|、]','',li):
                    print(line_[start:start + 20],'在本次生成的文案中出现过一次以上')
                    print_split_line()
                    invalide = True
                    break

        for bw in blocked_words:
            if bw in line:
                print(bw,'是屏蔽词')
                print_split_line()
                invalide = True
                break

        if not invalide:
            result.append(line)

    return '\n'.join(result)


# 开始仿写话术
def startGenScriptTask(scene_idx):
    global task_config,used_files,hist_script_content
    
    product_prompt = ''

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'])):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform']))

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'])):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product']))

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库')):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库'))

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'])):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer']))

    # 读取指令
    with open(os.path.join('config','指令文件',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'] + '.txt'),'r+',encoding = 'utf-8') as f:
        product_prompt = f.read()

        if re.findall('【屏蔽词】(.+)', product_prompt):
            blocked_words = re.findall('【屏蔽词】(.+)', product_prompt)[0].split('、')
        else:
            blocked_words = []

        product_prompt_ = re.sub('【屏蔽词】.+','',product_prompt)

    for i in range(5000000):
        if len(os.listdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer']))) > int(task_config[scene_idx]['cache_count']):
            print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'库存超过%d，暂停60秒' % int(task_config[scene_idx]['cache_count']))
            print_split_line()
            time.sleep(60)
            continue

        # 暂停生成
        if not task_config[scene_idx]['startGenScriptTasking']:
            break

        result = doubao(random_prompt(product_prompt_),temperature = None)

        print('生成的结果：',result)

        print_split_line()

        result_clean = clear_script(result,blocked_words,task_config[scene_idx]['live_platform'])

        print('清洗后的结果：',result_clean)

        print_split_line()

        if result_clean and len(result) > 2500:
            path = os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'],'%s.txt' % datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")) 
            
            with open(path,'w',encoding = 'utf-8') as f:
                f.write(result_clean)

            # 把新生成的文案加入到历史文案中去，为了和后续生成的文案进行去重
            hist_script_content[task_config[scene_idx]['live_platform']] += re.sub('[，|。|！|!|？|?|:|：|\n|-|——|;|；|、]','',result_clean) 

        # 这里睡眠一下为的是让其他任务有机会可以执行
        time.sleep(1)

    print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'暂停话术生成')
    print_split_line()


@app.route('/')
def login():
    return render_template('login.html')


@app.route('/index')
def index():
    return render_template('index.html', products = list(set(products)),streamers = streamers,platforms = platforms)


# 开始生成话术
@app.route('/startGenScript',methods = ['POST'])
def startGenScript():
    global task_config

    scene_idx = request.form.get('scene_idx')

    if not task_config[scene_idx]['live_platform'] or not task_config[scene_idx]['current_product'] or not task_config[scene_idx]['current_streamer']:
        return jsonify({'success':True,'message':''})

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'])):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform']))

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'])):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product']))

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库')):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库'))

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'])):
        os.mkdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer']))

    if not os.path.exists(os.path.join('config','指令文件',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'] + '.txt')):
        return jsonify({'success':False,'message':'请先部署好指令文件'})

    task_config[scene_idx]['startGenScriptTasking'] = True

    task_thread = threading.Thread(target = startGenScriptTask,args = (scene_idx,))
    task_thread.start()

    return jsonify({'success':True,'message':''})


# 暂停生成话术
@app.route('/stopGenScript',methods = ['POST'])
def stopGenScript():
    global task_config

    scene_idx = request.form.get('scene_idx')

    task_config[scene_idx]['startGenScriptTasking'] = False

    return jsonify({'success':True,'message':''})


# 开始下载话术
@app.route('/startDownloadScript',methods = ['POST'])
def startDownloadScript():
    global task_config

    product = request.form.get('product')
    streamer = request.form.get('streamer')
    platform = request.form.get('platform')

    if not product or not streamer or not platform:
        return jsonify({'success':True,'content':''})

    if not os.path.exists(os.path.join('config','话术脚本',platform)):
        return jsonify({'success':True,'platform':''})

    if not os.path.exists(os.path.join('config','话术脚本',platform,product)):
        return jsonify({'success':True,'content':''})

    if not os.path.exists(os.path.join('config','话术脚本',platform,product,'文本库')):
        return jsonify({'success':True,'content':''})

    if not os.path.exists(os.path.join('config','话术脚本',platform,product,'文本库',streamer)):
        return jsonify({'success':True,'content':''})

    with lock:
        files = [file for file in os.listdir(os.path.join('config','话术脚本',platform,product,'文本库',streamer)) if file.endswith('.txt')]

        if files:
            file_path = os.path.join('config','话术脚本',platform,product,'文本库',streamer,files.pop(0))

            with open(file_path,'r',encoding = 'utf-8') as f:
                content = f.read().strip()
            
            shutil.move(file_path,os.path.join('D:\\','历史文案',platform,'%s.txt' % ''.join(random.sample('abcdefghijklmnopqrstuvwxyz',8))))

            return jsonify({'success':True,'content':content})

        else:
            return jsonify({'success':True,'content':''})


# 清空话术
@app.route('/clearScript',methods = ['POST'])
def clearScript():
    global task_config

    scene_idx = request.form.get('scene_idx')

    # 先清空历史脚本
    for file in os.listdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'])):
        os.remove(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'],file))

    return jsonify({'success':True,'message':''})


# 选择产品
@app.route('/selectProduct',methods = ['POST'])
def selectProduct():
    global task_config

    scene_idx = request.form.get('scene_idx')
    current_product = request.form.get('product')

    task_config[scene_idx]['current_product'] = current_product

    return jsonify({'success':True,'message':'已切换到' + current_product})


# 选择音色
@app.route('/selectStreamer',methods = ['POST'])
def selectStreamer():
    global task_config

    scene_idx = request.form.get('scene_idx')
    current_streamer = request.form.get('streamer')

    task_config[scene_idx]['current_streamer'] = current_streamer

    return jsonify({'success':True,'message':'已切换到' + current_streamer})


# 选择缓存数
@app.route('/selectCacheCount',methods = ['POST'])
def selectCacheCount():
    global task_config

    scene_idx = request.form.get('scene_idx')
    cache_count = request.form.get('cache_count')

    task_config[scene_idx]['cache_count'] = cache_count or '100'

    return jsonify({'success':True,'message':''})


# 选择开播平台
@app.route('/selectLivePlatform',methods = ['POST'])
def selectLivePlatform():
    global task_config

    scene_idx = request.form.get('scene_idx')
    live_platform = request.form.get('live_platform')

    task_config[scene_idx]['live_platform'] = live_platform

    return jsonify({'success':True,'message':''})


# 选择场景
@app.route('/selectScene',methods = ['POST'])
def selectScene():
    global task_config

    scene_idx = request.form.get('scene_idx')
    
    task_config[scene_idx]['active'] = True


    return jsonify({'success':True,'message':''})


# 同步配置信息
@app.route('/update_task_config',methods = ['POST'])
def update_task_config():
    global task_config

    return jsonify(task_config)


# 查询话术库存
@app.route('/getScriptCount',methods = ['POST'])
def getScriptCount():
    global task_config

    scene_idx = request.form.get('scene_idx')

    if not task_config[scene_idx]['live_platform']:
        return jsonify({'success':True,'count': -1})

    if not task_config[scene_idx]['current_product']:
        return jsonify({'success':True,'count': -1})

    if not task_config[scene_idx]['current_streamer']:
        return jsonify({'success':True,'count': -1})

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'])):
        return jsonify({'success':True,'count': -1})

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'])):
        return jsonify({'success':True,'count': -1})

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库')):
        return jsonify({'success':True,'count': -1})

    if not os.path.exists(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'])):
        return jsonify({'success':True,'count': -1})

    data = [file for file in os.listdir(os.path.join('config','话术脚本',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],'文本库',task_config[scene_idx]['current_streamer'])) if file.endswith('.txt')]

    return jsonify({'success':True,'count':len(data)})


if __name__ == '__main__':
    app.run(debug = True,port = 6868,host = '0.0.0.0')
